import paramiko
import glog
import re

from setting import *

remote_dir = "./"


class SudoClient(paramiko.SSHClient):
    def sudo_exec_command(self, cmd: str):
        stdin, stdout, stderr = self.exec_command(f'sudo -S bash -c "{cmd}"')
        stdin.write(host_password + '\n')
        stdin.flush()
        out = stdout.read().decode()
        err = stderr.read().decode()
        print(out, err)
        return out,err
    
def create_clients():
    clients = []
    for i,ip in hosts_manage_ip.items():
        client = SudoClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            glog.info("Connecting to %s"%ip)
            client.connect(ip, ssh_port, host_username, host_password)
            clients.append((i, ip, client))
        except Exception:
            glog.fatal(f'cannot connect to {ip}(H{i})')
    return clients

def exit(clients):
    for i, ip, client in clients:
        print(f'H{i}:{ip} exiting.')
        client.close()

def set_data_plane_ip(clients):
    for i,ip,client in clients:
        print(f'========================H{i}==============================')
        cmd='ifconfig | grep ens -A 1'
        out,_=client.sudo_exec_command(cmd)
        for line in out.split('--'):
            port_name = re.search(r'ens[0-9]+', line)
            inet = re.search(r'inet ', line)
            if(port_name != None and inet == None):
                port_name = port_name.group()
                data_plane_ip = f'{host_data_ip_addr_prefix}{i}'
                cmd = f'ip addr add {data_plane_ip}/24 dev {port_name}\n'
                cmd += f'ifconfig | grep {port_name} -A 1\n'
                out,_ = client.sudo_exec_command(cmd)
                exist=re.search(r'inet ', out)
                if(exist == None):
                    print(f'Failed to set data plane ip on H{i}({ip}).')
                else:
                    print(f'Successed to set data plane ip on H{i}({ip}).')

def set_ipv6(clients):
    for i,ip,client in clients:
        print(f'========================H{i}==============================')
        cmd='lspci -nnk | grep -i Ether -A2'
        out,_=client.sudo_exec_command(cmd)
        for line in out.split('--'):
            pci = re.search(r'[0-9]+:([0-9a-z]+)\.[0-9]', line) 
            if(pci != None and ("virtio-pci" in line or "ixgbe" in line)):
                port_name = f'ens{int(pci.group(1), 16)}'
                data_plane_ip = f'{host_data_ip_addr_prefix}{i:x}'
                cmd = f'ip -6 addr add {data_plane_ip}/{prefix_len} dev {port_name}\n'
                cmd += f'sudo ip link set dev {port_name} mtu 1300\n'
                # cmd += f'ifconfig | grep {port_name} -A 1\n'
                print(cmd)
                out,_ = client.sudo_exec_command(cmd)

def set_mtu(clients):
    for i,ip,client in clients:
        print(f'========================H{i}==============================')
        cmd='lspci -nnk | grep -i Ether -A2'
        out,_=client.sudo_exec_command(cmd)
        for line in out.split('--'):
            pci = re.search(r'[0-9]+:([0-9a-z]+)\.[0-9]', line) 
            if(pci != None and ("virtio-pci" in line or "ixgbe" in line)):
                port_name = f'ens{int(pci.group(1), 16)}'
                data_plane_ip = f'{host_data_ip_addr_prefix}{i:x}'
                cmd = f'sudo ip link set dev {port_name} mtu 1300\n'
                # cmd += f'ifconfig | grep {port_name} -A 1\n'
                print(cmd)
                out,_ = client.sudo_exec_command(cmd)

def del_ipv6(clients):
    for i,ip,client in clients:
        print(f'========================H{i}==============================')
        cmd='lspci -nnk | grep -i Ether -A2'
        out,_=client.sudo_exec_command(cmd)
        for line in out.split('--'):
            pci = re.search(r'[0-9]+:([0-9a-z]+)\.[0-9]', line) 
            if(pci != None and ("virtio-pci" in line or "ixgbe" in line)):
                port_name = f'ens{int(pci.group(1), 16)}'
                data_plane_ip = f'{host_data_ip_addr_prefix}{i:x}'
                cmd = f'ip -6 addr del {data_plane_ip}/{prefix_len} dev {port_name}\n'
                # cmd += f'ifconfig | grep {port_name} -A 1\n'
                print(cmd)
                out,_ = client.sudo_exec_command(cmd)

def input_command(clients):
    command = input()
    for i, ip, client in clients:
        print(f'========================H{i}==============================')
        client.sudo_exec_command(command)


def reboot(clients):
    for i,ip,client in clients:
        print(f'========================H{i}==============================')
        cmd = 'reboot now'
        client.sudo_exec_command(cmd)

def set_hostname(clients):
    for i,ip,client in clients:
        hostname = f'h{i}'
        cmd = f'echo {hostname} > /etc/hostname'
        print(cmd)
        client.sudo_exec_command(cmd)

def main():
    # 操作选项和对应命令
    options = [
        set_data_plane_ip,
        set_ipv6,
        set_mtu,
        del_ipv6,
        input_command,
        set_hostname,
        reboot,
        exit
    ]
    glog.info("Connecting to remote hosts...")
    clients = create_clients()
    stop = False
    hint = ""
    for i, option in zip(range(len(options)),options):
        hint += '[%d] %s\n'%(i, option.__name__)
    try:
        while not stop:
            id = int(input("请输入操作选项的编号:\n"+hint))
            options[id](clients)
            if id == len(options)-1:
                stop = True
    finally:
        exit(clients)

if __name__ == '__main__':
    main()

            